"""

【注意：开发中，还没有完成。】

"""

import datetime
import gradio as gr
from PyCmpltrtok.common import uuid, to_decode, to_int
import time
from redis_interface import text_gen
import redis
from http.cookies import SimpleCookie

# 连接本机环境变量R_PORT指明的端口上的redis的第0个数据库
rdbr = redis.Redis('127.0.0.1', int(os.environ['R_PORT']), 0)
rdbr.get('try_it')

CHAR_LIMIT = 450
PROMPT = '（因平台限制，内容未完全展示，输入句号（.或。）来继续展示。）'
GO_ON_SET = set(['.', '。'])
WAIT_LAST_TURN = '（请等待上一轮对话完成。）'
THRESH_S = 20  # 单例锁最多等待这么多秒


def chat_once(xinput, username, xlog):
    """输入文本，用模型返回输出。"""    
    if username is None:
        # 为了健壮性
        return '', []
    
    # 单例限制
    if xinput not in GO_ON_SET:
        xis_in = rdbr.sismember('users_now', username)
        ts_now = time.time_ns()
        if xis_in:
            ts = rdbr.hget('username2ts', username)
            ts = to_int(ts)
            if ts and ts_now - ts < THRESH_S * 1e9:
                xlog.append([
                    xinput,
                    WAIT_LAST_TURN
                ])
                return '', xlog
        rdbr.hset('username2ts', username, ts_now)
        rdbr.sadd('users_now', username)

    def get_result():
        """
        闭包函数，拿数据
        """
        # 最后结果
        xoutput = rdbr.hget('username2output', username)
        
        # 当前偏移
        xoffset = rdbr.hget('username2offset', username)
        xoffset = to_int(xoffset)
        xoffset_ = xoffset
        # 分推理结束和未结束两种
        if xoutput is not None:  # 结束
            xoutput = xoutput.decode('utf8')
            
            # 取本次输出
            xoutput = xoutput[xoffset:]
            xlen_output = len(xoutput)
            
            # 分输出完成未完成两种
            if xlen_output <= CHAR_LIMIT:
                # 输出完成，清理username2dynamic以标志输出完成
                rdbr.hdel('username2dynamic', username)
                rdbr.hdel('username2offset', username)
                return xoutput
            else:
                # 输出未完成，分段输出
                xres = xoutput[:CHAR_LIMIT]
                
        else:  # 未结束
            
            # 从动态响应里拿
            xres = rdbr.hget('username2dynamic', username)
            xres = to_decode(xres)
            
            # 限制一下长度
            xres = xres[xoffset:xoffset+CHAR_LIMIT]
            
        print('output:', xres)
        # 记录偏移
        xlen = len(xres)
        rdbr.hset('username2offset', username, xoffset_ + xlen)
        
        return xres + PROMPT
    
    if xinput in GO_ON_SET:  # "继续显示"
        # 从动态响应里拿
        xres = rdbr.hget('username2dynamic', username)
        # 没有动态响应，说明输出结束
        if xres is None:
            return '（已经没有需要继续输出的内容了。）'
        
        print('GO ON >>>>')
        return get_result()
    
    # 生成
    text_gen(rdbr, xinput, username)
    
    # 既然限制5秒，干脆等3.75秒，让大模型去生成
    time.sleep(3.75)
    
    # 3.75秒后，拿结果
    xres = get_result()
    return xres

    

    # 返回空字符串清空输入文本框，返回聊天历史
    return '', xlog


def embed_user_id_in_state_comp(request: gr.Request):
    """
    页面加载的事件handler

    从HTTP request中获取cookie来确定用户名，如果没有则生成一个
    :param request: HTTP request data
    :return: 将用户名发送给浏览器中存放用户名的组件
    """
    xuuid = uuid()

    #从cookie中获取username
    username = None
    try:
        cookie_str = request.headers.cookie
        print('cookie_str:', cookie_str)
        cookie = SimpleCookie()
        cookie.load(cookie_str)
        username = cookie.get('username', None)
        print('user:', username)
        username = username.value
    except Exception:
        try:
            username = request.cookies['username']
        except Exception:
            pass

    # 获取不到，新建一个username
    if username is None:
        print('new user!')
        username = f'u{xuuid}'

    # print username
    print('user:', username)

    xlog = get_history(mdb, username, more_info=True)

    return username, username, username, xlog


# 用户名写入cookie的JS
js_username2cookie = """
    async () => {{
            setTimeout(() => {{
                var xnew = document.querySelector('#username.prose').innerHTML;
                console.log('username', xnew)
                if ('' == xnew) {
                    return;
                }
                document.cookie = 'username=' + xnew + '; max-age=' + 3600*24*30*6;
            }}, 1);
        }}
"""


def do_the_init():
    """
        点击”开始聊天“的事件的handler
        :return: 隐藏”开始聊天“，显示其他组件。
    """
    return (
        gr.update(visible=False),
        *([gr.update(visible=True)] * 3),
    )
    
    
def show_text_out_cache(username):
    if username is None:
        # 为了健壮性
        return ''

    cache = rdb.hget('username2dynamic', username)
    if cache is None:
        return ''
    cache = cache.decode('utf8')
    return cache


# 构建Gradio界面
with gr.Blocks(analytics_enabled=False) as demo:

    # 用户名相关控件
    cmp_username_state = gr.State()
    cmp_username_html = gr.HTML(visible=False, elem_id='username')
    cmp_username_display = gr.Textbox(interactive=False, label='用户名')

    # 启动按钮
    cmp_start_btn = gr.Button('开始聊天')

    # 聊天历史
    cmp_chatbot = gr.Chatbot().style(height=350)
    
    # 动态响应输出控件
    cmp_text_out_cache = gr.Textbox(interactive=False, label='动态响应', visible=False, )
    
    # 输入文本框
    cmp_input = gr.Textbox(interactive=True, label='输入', visible=False, )

    # 输入按钮
    cmp_input_button = gr.Button('发送', visible=False, )


    # 输入按钮点击事件
    cmp_input_button.click(chat_once, [cmp_input, cmp_username_state], [
        cmp_input,
        cmp_chatbot,
    ], queue=False)

    # 加载时获取或设置cookie，并展示聊天历史
    demo.load(embed_user_id_in_state_comp, None, [
        cmp_username_state,
        cmp_username_display,
        cmp_username_html,
        cmp_chatbot,
    ], queue=False)
    # cmp_username_html发生变化时，触发JS，把生成的username放入页面cookie
    cmp_username_html.change(None, None, None, _js=js_username2cookie, queue=False)
    
    # 启动
    cmp_start_btn.click(do_the_init, None, [
        cmp_start_btn,
        cmp_text_out_cache,
        cmp_input,
        cmp_input_button,
    ], queue=False)
    
    # 输出动态响应
    cmp_start_btn.click(show_text_out_cache, cmp_username_state, cmp_text_out_cache, every=0.1, )

# 使用queue（目前这个例子太简单，没有用到queue，后面会用到。）
demo.queue(
    concurrency_count=10,
    status_update_rate='auto',
    # status_update_rate=0.02,
)

# 启动Gradio
# demo.launch(server_name='0.0.0.0', server_port=6006, share=True, debug=True)  # 带内网穿透
demo.launch(server_name='0.0.0.0', server_port=6006, debug=True)
